#!/usr/bin/env python 
# -*- coding:utf-8 -*-

import os
import cartopy.crs as ccrs
import matplotlib.colors as colors
import matplotlib.pyplot as plt
import numpy as np
import shapefile
from matplotlib.patches import PathPatch
from matplotlib.path import Path
from scipy.interpolate import Rbf

# ===================================参数====================================
font = "SimSun"
lat = [28.0, 28.0, 28.0, 28.0, 28.0, 28.0, 28.0, 29.0, 29.0, 29.0, 29.0, 29.0, 29.0, 29.0, 30.0, 30.0, 30.0, 30.0, 30.0, 30.0, 30.0, 31.0, 31.0, 31.0, 31.0, 31.0, 31.0, 31.0, 32.0, 32.0, 32.0, 32.0, 32.0, 32.0, 32.0, 33.0, 33.0, 33.0, 33.0, 33.0, 33.0, 33.0]
lon = [105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0, 105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0, 105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0, 105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0, 105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0, 105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0]
val = [2.4, 2.3, 2.6, 2.5, 2.5, 2.3, 2.1, 3.3, 4.3, 3.6, 4.4, 3.4, 4.1, 4.4, 4.9, 5.3, 5.6, 4.7, 4.5, 4.6, 7.5, 5.2, 4.9, 5.6, 5.8, 4.9, 5.5, 6.2, 5.0, 5.6, 5.6, 5.1, 6.5, 5.9, 6.4, 5.7, 7.0, 7.4, 7.1, 7.0, 6.4, 7.7]
min_lon = 105.0
max_lon = 111.0
min_lat = 28.0
max_lat = 33.0
shp_name = "/data/shp/group/chongqing/cq_City.shp"
levels = [2.1, 2.9, 3.7, 4.5, 5.3, 6.1, 6.9, 7.7]
color_map = ["#244ba6", "#145eae", "#64b7f8", "#a0f9f4", "#fefebe", "#fdd283", "#f88b51", "#dd3d2d", "#9d2539"]
shp_encoding = "gbk"
save_path = "./"
image_name = "test.png"
# ===========================================================================


# 用来正常显示中文 SimHei
plt.rcParams['font.sans-serif'] = [font]
# 用来正常显示负号
plt.rcParams['axes.unicode_minus'] = False

# 插值为10公里的格点经纬度长度
lon_length = int((max_lon - min_lon) * 10) + 1
lat_length = int((max_lat - min_lat) * 10) + 1

# 经纬度处理
olon = np.linspace(min_lon, max_lon, lon_length)
olat = np.linspace(min_lat, max_lat, lat_length)

# 转成地理经纬度
olon, olat = np.meshgrid(olon, olat)

# 插值处理
func = Rbf(lon, lat, val, function='linear')
rhu_data_new = func(olon, olat)

print("插值后的格点数据")
print(rhu_data_new)

# 自定义色卡处理
under_color = color_map[0]
over_color = color_map[len(color_map) - 1]
my_cmap = colors.ListedColormap(color_map)
my_cmap.set_under(under_color)
my_cmap.set_over(over_color)
norm3 = colors.BoundaryNorm(levels, my_cmap.N)

# 投影方式
projection = ccrs.Mercator()

# 创建图实例 创建设置像素
fig = plt.figure(dpi=300)
ax = fig.add_subplot(111, projection=projection)

# 设定要显示的边界区域
ax.set_extent([min_lon, max_lon, min_lat, max_lat], crs=projection)

# if color_map_dto.show_county_boundary:
#     # 如果是重庆不会进这个if，只有省+市才会进来
#     if color_map_dto.city_shp_path:
#         # ====================开始叠加市信息=====================
#         ax.add_geometries(Reader(color_map_dto.city_shp_path).geometries(), ccrs.Mercator(),
#                           facecolor='none',
#                           edgecolor='#000000', linewidth=0.8)
#         # ====================叠加市信息结束=====================
#     # ====================开始叠加区县信息=====================
#     ax.add_geometries(Reader(color_map_dto.county_boundary_shp).geometries(), ccrs.Mercator(),
#                       facecolor='none',
#                       edgecolor='#000000', linewidth=0.1, alpha=0.2)
#     # 区域外边界
#     ax.add_geometries(Reader(shp_name).geometries(), ccrs.Mercator(), facecolor='none',
#                       edgecolor='#00A2D5', linewidth=1.2)
#     # ====================叠加区县信息结束=====================
#
# if color_map_dto.show_river:
#     # ====================开始叠加河流信息=====================
#     ax.add_geometries(Reader(color_map_dto.river_shp).geometries(), ccrs.Mercator(),
#                       facecolor='none',
#                       edgecolor='#0063BF', linewidth=0.8)
#     # ====================叠加河流信息结束=====================
#
# if color_map_dto.show_text:
#     # ====================开始叠加站号信息=====================
#     text_list = color_map_dto.text_list
#     text_lat_list = color_map_dto.text_lat_list
#     text_lon_list = color_map_dto.text_lon_list
#     for i in range(len(text_list)):
#         ax.text(text_lon_list[i], text_lat_list[i], text_list[i], fontsize=5, color="black")
#     # ====================叠加站号信息结束=====================

# 等值面处理
cs = ax.contourf(olon, olat, rhu_data_new, levels=levels, cmap=my_cmap, norm=norm3, extend='both',
                 transform=projection, alpha=1)

# 读取shp文件，获取要裁剪的区域clip(环流图不用裁剪)
if shp_name:
    sf = shapefile.Reader(shp_name, encoding = shp_encoding)
    for shp_record in sf.shapeRecords():
        vertices = []
        codes = []
        pts = shp_record.shape.points
        prt = list(shp_record.shape.parts) + [len(pts)]
        for i in range(len(prt) - 1):
            for j in range(prt[i], prt[i + 1]):
                vertices.append((pts[j][0], pts[j][1]))
            codes += [Path.MOVETO]
            codes += [Path.LINETO] * (prt[i + 1] - prt[i] - 2)
            codes += [Path.CLOSEPOLY]
        clip = Path(vertices, codes)
        clip = PathPatch(clip, transform=ax.transData)

    # 裁剪显示区域
    for contour in cs.collections:
        contour.set_clip_path(clip)

# cartopy 设置背景透明及去除边框
ax.outline_patch.set_visible(False)
ax.background_patch.set_visible(False)
# 判断目录是否存在，不存在则创建
if os.path.exists(save_path) is False:
    os.makedirs(save_path)
# 保存图片，tight让保存的图片更加紧促
plt.savefig(save_path + image_name, bbox_inches='tight', transparent=True, pad_inches=0)
plt.clf()
plt.close(fig)